//! Conversion down from CIR to lang IR.

use std::borrow::Borrow;

use stencil_codegen::Error;
use stencil_mir::{
    AbortReason, AssignTarget, ContainerData, CtorAssignments, MirBlock, MirFunction, MirModule,
    MirSignature, MirStatement, TypedExpr,
};
use stencil_tysys::{
    ArgRef, Endianness, EnvExpr, FieldRef, FnContext, FnContextExt, IntType, LiteralValue,
    LocalRef, ValueType,
};
use stencil_util::IdentFormat;

use crate::ir::{Lvalue, LvalueRef};

use super::context::*;
use super::ir::{
    ClassFieldDef, Exception, FnArg, FnSignature, Identifier, PyBlock, PyClassDef, PyDef, PyExpr,
    PyFnDef, PyModule, PyStatement, PyType, PyValue,
};

pub fn conv_module(m: &MirModule) -> Result<PyModule, Error> {
    let mut defs = Vec::new();

    for cont in m.containers() {
        let mut cls_field_defs = Vec::new();
        let mut cls_fn_defs = Vec::new();

        // Generate the fields.
        for f in cont.data().fields() {
            let ident = Identifier::from_field_name(f.name());
            let ty = PyType::from_mir(f.ty());
            cls_field_defs.push(ClassFieldDef::new(ident, ty));
        }

        // Generate the functions.
        let fun_context = FunctionContext {
            m,
            cont: Some(cont.data()),
        };

        for fun_entry in cont.functions() {
            let def = conv_function(fun_entry.imp(), &fun_context)?;
            cls_fn_defs.push(def);
        }

        // Assemble the final classdef.
        let cls_def = PyClassDef {
            name: Identifier::from_type_name(cont.name()),
            doc: None,
            is_dataclass: false, // FIXME better decide when we need to use this
            fields: cls_field_defs,
            functions: cls_fn_defs,
        };

        defs.push(PyDef::Class(cls_def));
    }

    Ok(PyModule {
        initial_comment: None,
        defs,
    })
}

/// Context for a function.
struct FunctionContext<'c> {
    m: &'c MirModule,
    cont: Option<&'c ContainerData>,
}

impl<'c> FunctionContext<'c> {
    /// Gets the identifier of a field ref.
    pub fn get_field_ident(&self, fr: FieldRef) -> Result<Identifier, Error> {
        match self.cont {
            Some(cd) => {
                let field = cd.get_field(fr).ok_or(Error::UnknownField(fr))?;
                Ok(Identifier::from_field_name(field.name()))
            }
            None => Err(Error::NotContainer),
        }
    }

    /// Converts a MIR assignment target to a Python lvalue.
    pub fn conv_assign_target(&self, target: AssignTarget) -> Result<Lvalue, Error> {
        match target {
            AssignTarget::Local(local) => {
                let ident = gen_local_ident(local);
                Ok(Lvalue::Ref(LvalueRef::Local(ident)))
            }
            AssignTarget::InstanceField(field) => {
                let field_ident = self.get_field_ident(field)?;
                Ok(Lvalue::RefField(LvalueRef::ClassSelf, field_ident))
            }
        }
    }

    /// Gets a Python expression for an env.
    pub fn get_env_expr(&self, env_expr: &EnvExpr) -> Result<PyExpr, Error> {
        match env_expr {
            EnvExpr::Field(fr) => Ok(PyExpr::new_self_field(self.get_field_ident(*fr)?)),
            EnvExpr::Arg(ar) => Ok(PyExpr::Ident(gen_fn_arg_ident(*ar))),
            EnvExpr::Local(lr) => Ok(PyExpr::Ident(gen_local_ident(*lr))),
        }
    }
}

/// Generates a safe name for a local given a local ref.
fn gen_local_ident(lr: LocalRef) -> Identifier {
    let i = lr.inner() as u16;
    Identifier::try_from(format!("local{i}")).unwrap()
}

// TODO figure out how to override arg names
fn gen_fn_arg_ident(ar: ArgRef) -> Identifier {
    let i = ar.inner() as u16;
    Identifier::try_from(format!("arg{i}")).unwrap()
}

/// Converts a function, which might be a constructor.
fn conv_function(fun: &MirFunction, ctx: &FunctionContext<'_>) -> Result<PyFnDef, Error> {
    let is_ctor = fun.sig().is_constructor();

    let name = if is_ctor {
        Identifier::try_from(format!("__str__")).unwrap()
    } else {
        Identifier::from_fn_name(fun.name())
    };

    let def = PyFnDef {
        name,
        doc: None,
        is_constructor: is_ctor,
        is_static: !fun.sig().is_instance(),
        sig: conv_signature(fun.sig(), ctx)?,
        body: conv_block(fun.root_block(), ctx)?,
    };

    Ok(def)
}

/// Converts a function signature.
fn conv_signature(sig: &MirSignature, _ctx: &FunctionContext<'_>) -> Result<FnSignature, Error> {
    let mut args = Vec::new();

    for (i, a) in sig.arg_tys().iter().enumerate() {
        let ident = gen_fn_arg_ident(ArgRef::from(i as u16));
        args.push(FnArg::new(ident, Some(PyType::from_mir(a)), None))
    }

    let ret_ty = sig.ret_ty().map(PyType::from_mir);

    Ok(FnSignature::new(sig.is_instance(), args, ret_ty))
}

fn conv_block(block: &MirBlock, ctx: &FunctionContext<'_>) -> Result<PyBlock, Error> {
    let mut statements = Vec::new();

    for stmt in block.statements() {
        // Filter out statements we have to do special handling for.
        match stmt {
            // Python uses weird scoping so we don't do this.
            MirStatement::DeclareHint(_, _) => continue,

            // Constructor is just multiple statements.
            MirStatement::Construct(ctor) => {
                conv_ctor_into(ctor, ctx, &mut statements)?;
                continue;
            }

            _ => {}
        }

        let py_stmt = conv_statement(stmt, ctx)?;
        statements.push(py_stmt);
    }

    Ok(PyBlock::new(statements))
}

fn conv_statement(stmt: &MirStatement, ctx: &FunctionContext<'_>) -> Result<PyStatement, Error> {
    let s = match stmt {
        MirStatement::Comment(msg) => PyStatement::Comment(msg.clone()),

        MirStatement::Assign(target, expr) => {
            let lv = ctx.conv_assign_target(*target)?;
            let py_expr = conv_expr(expr, ctx)?;
            PyStatement::Assign(lv, py_expr)
        }

        MirStatement::AssignIndex(target, index, value) => {
            let lv = ctx.conv_assign_target(*target)?;
            let idx_pe = conv_expr(index, ctx)?;
            let val_pe = conv_expr(value, ctx)?;
            PyStatement::AssignIndex(lv, idx_pe, val_pe)
        }

        MirStatement::AssignBufInt(target, offset, value, ienc) => {
            let lv = ctx.conv_assign_target(*target)?;
            let off_pe = conv_expr(offset, ctx)?;
            let val_pe = conv_expr(value, ctx)?;
            PyStatement::AssignBufInt(lv, off_pe, val_pe, *ienc)
        }

        MirStatement::ReturnVoid => PyStatement::ReturnNone,
        MirStatement::ReturnValue(expr) => PyStatement::Return(conv_expr(expr, ctx)?),

        MirStatement::If(cond, block) => {
            let cond_pe = conv_expr(cond, ctx)?;
            let pb = conv_block(block, ctx)?;
            PyStatement::If(cond_pe, Box::new(pb))
        }

        MirStatement::IfElse(cond, then_block, else_block) => {
            let cond_pe = conv_expr(cond, ctx)?;
            let ptb = conv_block(then_block, ctx)?;
            let peb = conv_block(else_block, ctx)?;
            PyStatement::IfElse(cond_pe, Box::new(ptb), Box::new(peb))
        }

        // The rest we don't implement yet because we don't generate them yet, or we catch them before.
        s => {
            eprintln!("error (python): unsupported statement {s:?}");
            return Err(Error::Unimplemented);
        }
    };

    Ok(s)
}

fn conv_ctor_into(
    ctor: &CtorAssignments,
    ctx: &FunctionContext<'_>,
    stmts: &mut Vec<PyStatement>,
) -> Result<(), Error> {
    // Add a comment to clarify what we're doing here.
    stmts.push(PyStatement::Comment("setting the fields".to_owned()));

    for assignment in ctor.assignments() {
        // We are essentially assembling regular assignments together on the fly.
        let expr = conv_expr(assignment.expr(), ctx)?;
        let assign_target = AssignTarget::InstanceField(assignment.field());
        let stmt = PyStatement::Assign(ctx.conv_assign_target(assign_target)?, expr);
        stmts.push(stmt);
    }

    Ok(())
}

fn conv_expr(expr: &TypedExpr, ctx: &FunctionContext<'_>) -> Result<PyExpr, Error> {
    use stencil_mir::TaggedExpr;

    match expr {
        TaggedExpr::Const(_, lv) => conv_literal(lv, ctx),
        TaggedExpr::Env(_, e) => ctx.get_env_expr(e),

        TaggedExpr::Monop(_, op, e) => {
            let pe = conv_expr(e, ctx)?;
            Ok(PyExpr::Monop(*op, Box::new(pe)))
        }

        TaggedExpr::Binop(_, op, lhs, rhs) => {
            let plhs = conv_expr(lhs, ctx)?;
            let prhs = conv_expr(rhs, ctx)?;
            Ok(PyExpr::Binop(*op, Box::new(plhs), Box::new(prhs)))
        }

        TaggedExpr::Cond(_, cond, then_expr, else_expr) => {
            let pc = conv_expr(cond, ctx)?;
            let pt = conv_expr(then_expr, ctx)?;
            let pe = conv_expr(else_expr, ctx)?;
            Ok(PyExpr::Ternary(Box::new(pc), Box::new(pt), Box::new(pe)))
        }

        TaggedExpr::Cast(_, target_ty, e) => {
            let pty = PyType::from_mir(target_ty); // TODO figure out how to insert sanity checks
            let pe = conv_expr(e, ctx)?;
            Ok(PyExpr::Cast(pty, Box::new(pe)))
        }

        TaggedExpr::BufferLen(_, e) => {
            let pe = conv_expr(e, ctx)?;
            Ok(PyExpr::Len(Box::new(pe)))
        }

        TaggedExpr::BufferRead(_, target, start, end) => {
            let pt = conv_expr(target, ctx)?;
            let ps = conv_expr(start, ctx)?;
            let pe = conv_expr(end, ctx)?;
            Ok(PyExpr::BytebufReadSlice(
                Box::new(pt),
                Box::new(ps),
                Box::new(pe),
            ))
        }

        TaggedExpr::BufferReadInt(_, target, offset, ienc) => {
            let pt = conv_expr(target, ctx)?;
            let po = conv_expr(offset, ctx)?;
            Ok(PyExpr::BytebufReadInt(Box::new(pt), Box::new(po), *ienc))
        }
    }
}

fn conv_literal(v: &LiteralValue, _ctx: &FunctionContext<'_>) -> Result<PyExpr, Error> {
    let l = match v {
        LiteralValue::Void => PyValue::None,
        LiteralValue::U8(v) => PyValue::Int(*v as i64),
        LiteralValue::I64(v) => PyValue::Int(*v as i64),
        LiteralValue::U64(v) => PyValue::Int(*v as i64),
        LiteralValue::USize(v) => PyValue::Int(*v as i64),
        LiteralValue::ISize(v) => PyValue::Int(*v as i64),
        LiteralValue::Bool(v) => PyValue::Bool(*v),
    };

    Ok(PyExpr::Literal(l))
}

/*
fn conv_dtype(dt: &ir::DataType, m: &ir::Module) -> Result<PyDef, Error> {
    let buf_field = ClassFieldDef::new(
        Identifier::try_from("buf".to_owned()).expect("codegen/python: buf field name"),
        super::ir::PyType::Bytearray,
    );

    let mut fns = Vec::new();

    // Contexts for regular member functions.
    let gen_ctx = GenContext::new("self.buf".to_owned());

    for f in dt.fn_entries() {
        fns.push(conv_function(f, dt, m, f.code(), &gen_ctx)?);
    }

    let cd = PyClassDef {
        name: Identifier::from_ident_string(dt.ident(), IdentFormat::UpperCamel),
        doc: None,
        is_dataclass: false,
        fields: vec![buf_field],
        functions: fns,
    };

    Ok(PyDef::Class(cd))
}


fn conv_function(
    fun: &ir::Function,
    dt: &ir::DataType,
    m: &ir::Module,
    tctx: &impl FnContext,
    gctx: &GenContext,
) -> Result<PyFnDef, Error> {
    let code = fun.code();

    // Construct the function signature.
    let mut sig = FnSignature::instance().with_rty(conv_ty(fun.return_ty())?);
    for i in 0..code.num_args() {
        let arg = ArgRef::from(i as u16);
        let arg_name = Identifier::try_from(gen_fn_arg_ident(i)).unwrap();
        let arg_ty = conv_ty(code.get_arg_ty(arg).unwrap())?;
        sig = sig.with_arg(FnArg::new_simple(arg_name).with_type(arg_ty));
    }

    // TODO handle constructor or something?

    let mut body = conv_block(code.root_block(), fun, m, tctx, gctx)?;

    if let Some(rv) = code.root_block().result() {
        let ret_stmt = Statement::Return(gctx.gen_arg_pyexpr(rv));
        body.statements.push(ret_stmt);
    }

    let fd = PyFnDef {
        name: Identifier::from_ident_string(fun.ident(), IdentFormat::Snake),
        doc: None,
        is_constructor: false,
        is_static: false,
        sig,
        body,
    };

    Ok(fd)
}

fn conv_block(
    block: &ssa::Block,
    fun: &ir::Function,
    m: &ir::Module,
    tctx: &impl FnContext,
    gctx: &GenContext,
) -> Result<Block, Error> {
    let mut statements = Vec::new();
    for ent in block.entries() {
        let stmt = conv_entry(ent, fun, m, tctx, gctx)?;
        statements.push(stmt);
    }

    // TODO somehow handle result

    Ok(Block { statements })
}

fn conv_entry(
    ent: &ssa::SsaEntry,
    fun: &ir::Function,
    m: &ir::Module,
    tctx: &impl FnContext,
    gctx: &GenContext,
) -> Result<Statement, Error> {
    match ent {
        // Specially handle conditionals.
        MirStatement::Assignment(target, MirExpr::Cond(ca, tb, eb)) => {
            // FIXME this is broken, it doesn't do the outputs of the blocks right
            let tblock = conv_block(tb, fun, m, tctx, gctx)?;
            let eblock = conv_block(eb, fun, m, tctx, gctx)?;
            Ok(Statement::IfElse(
                gctx.gen_arg_pyexpr(ca),
                Box::new(tblock),
                Box::new(eblock),
            ))
        }

        MirStatement::Assignment(target, expr) => {
            let tlv = gctx.gen_target_lvalue(target);
            let expr = conv_expr(expr, tctx, gctx)?;
            Ok(Statement::Assignment(tlv, expr))
        }

        MirStatement::AbortIf(cond, reason) => {
            let raise_stmt = Statement::Raise(conv_abort_reason(reason));
            let block = Block::single(raise_stmt);
            Ok(Statement::If(gctx.gen_arg_pyexpr(cond), Box::new(block)))
        }

        MirStatement::WriteBuf(_target, off, value_arg) => {
            let src_ty = tctx
                .resolve_ssa_arg_type(value_arg)
                .expect("codegen/python: bad arg desc");

            match src_ty {
                ValueType::ByteArray(_size) => {
                    eprintln!("error: byte array copies not supported yet");
                    Err(Error::Unimplemented)
                }

                ValueType::ByteSlice => {
                    eprintln!("error: byte slice copies not supported yet");
                    Err(Error::Unimplemented)
                }

                ty if ty.is_integer() => {
                    let ity = IntType::try_from(ty)
                        .expect("codegen/python: inconsistent type integeryness [ICE]");

                    // FIXME support the proper write lvalues
                    Ok(Statement::BufIntAssignment(
                        gctx.gen_buffer_lvalue(),
                        gctx.gen_arg_pyexpr(off),
                        gctx.gen_arg_pyexpr(value_arg),
                        ity,
                        Endianness::Big,
                    ))
                }

                _ => {
                    eprintln!("error (conv_entry): don't know what this is {src_ty:?}");
                    Err(Error::Unimplemented)
                }
            }
        }

        MirStatement::Return => Ok(Statement::ReturnNone),
        MirStatement::ReturnValue(rv) => Ok(Statement::Return(gctx.gen_arg_pyexpr(rv))),
    }
}

fn conv_expr(
    expr: &ssa::SsaExpr,
    _tctx: &impl FnContext,
    gctx: &GenContext,
) -> Result<PyExpr, Error> {
    use ssa::SsaExpr;
    let pexpr = match expr {
        SsaExpr::Literal(v) => PyExpr::Literal(PyValue::from_common_ir(*v)),
        SsaExpr::Copy(v) => gctx.gen_arg_pyexpr(v),
        SsaExpr::Monop(o, v) => PyExpr::Monop(*o, Box::new(gctx.gen_arg_pyexpr(v))),
        SsaExpr::Binop(o, l, r) => PyExpr::Binop(
            *o,
            Box::new(gctx.gen_arg_pyexpr(l)),
            Box::new(gctx.gen_arg_pyexpr(r)),
        ),
        SsaExpr::Cond(_, _, _) => unreachable!(),
        SsaExpr::Ternary(c, t, e) => PyExpr::Ternary(
            Box::new(gctx.gen_arg_pyexpr(c)),
            Box::new(gctx.gen_arg_pyexpr(t)),
            Box::new(gctx.gen_arg_pyexpr(e)),
        ),
        SsaExpr::Cast(t, v) => {
            PyExpr::Cast(PyType::from_common_ir(t), Box::new(gctx.gen_arg_pyexpr(v)))
        }
        SsaExpr::BufferLen(buf) => {
            let buf_pe = gctx.gen_arg_pyexpr(buf);
            PyExpr::Len(Box::new(buf_pe))
        }
        SsaExpr::ReadInt(ienc, buf, v) => {
            let buf_expr = gctx.gen_arg_pyexpr(buf);
            let off_expr = gctx.gen_arg_pyexpr(v);
            PyExpr::BytebufReadInt(Box::new(buf_expr), Box::new(off_expr), *ienc)
        }

        // TODO implement this properly
        SsaExpr::ReadBuf(_, _, _) => PyExpr::Literal(PyValue::None),
    };

    Ok(pexpr)
}*/

fn conv_abort_reason(reason: &AbortReason) -> Exception {
    let s = match reason {
        AbortReason::BufMissized(_missized) => "buffer missized",
        AbortReason::Malformed => "malformed",
        AbortReason::Oob => "index out of bounds",
    };
    Exception::ValueError(s.to_string())
}

fn conv_ty(ty: &ValueType) -> Result<PyType, Error> {
    if ty.is_integer() {
        return Ok(PyType::Int);
    }

    Ok(match ty {
        ValueType::Void => PyType::None,
        ValueType::Bool => PyType::Bool,
        ValueType::List(ty) => PyType::List(Box::new(conv_ty(&ty)?)),
        ValueType::ByteArray(_) => PyType::Bytearray,
        // TODO add "unsupported type" variant
        _ => {
            eprintln!("error (conv_ty): oh god I don't know what this is either {ty:?}");
            return Err(Error::Unimplemented);
        }
    })
}
